from .segformer import SegFormer
from .hrnet_ocr import hrnet18, hrnet48, hrnet32
from .deeplabv3plus import DeepLabv3plus
from .bisenetv2 import BiSeNetV2

__all__ = ['SegFormer',
           'hrnet18', 'hrnet48', 'hrnet32',
           'DeepLabv3plus',
           'BiSeNetV2']


def get_model(args):
    if args.model == "segformer":
        model = SegFormer(variant=args.variant, num_classes=args.num_classes)
        
    elif args.model == "deeplabv3plus":
        model = DeepLabv3plus(variant=args.variant, num_classes=args.num_classes, pretrained=args.pretrained)
    elif args.model == "hrnet18" or args.model == "hrnet48" or args.model == "hrnet32":
        model = eval(args.model)(num_classes=args.num_classes)
    elif args.model == "bisenetv2":
        model = BiSeNetV2(n_classes=args.num_classes)
    else:
        raise ("model register error")

    return model
